import baostock as baostock
from libs.Objects import Objects
import pandas as pd
from cachetools import cached, TTLCache

# 创建一个最大容量为 100，过期时间为 10 秒的缓存
cache_AllTradingDate = TTLCache(maxsize=999999, ttl=3600)
cache_StockBasic = TTLCache(maxsize=999999, ttl=9000)


class StockBase(object):
    loginInfo = None

    def __init__(self):
        pass

    @staticmethod
    def reConnect():
        StockBase.loginInfo = None
        return StockBase.getBaostock()

    @staticmethod
    def getBaostock():
        if not StockBase.loginInfo:
            StockBase.loginInfo = baostock
            StockBase.loginInfo.login()
        return StockBase.loginInfo

    @staticmethod
    def exit():
        if StockBase.loginInfo:
            StockBase.loginInfo.logout()
            StockBase.loginInfo = None

    @staticmethod
    @cached(cache_AllTradingDate)
    def getAllTradingDate():
        return StockBase().__loadTradeDate()

    @staticmethod
    def getTradingDays(startDay, endDay):
        data_list = StockBase.getAllTradingDate()
        ret = list(filter(lambda x: startDay <= x <= endDay, data_list))
        return sorted(ret)

    @staticmethod
    @cached(cache_StockBasic)
    def getStockBasic():
        return StockBase().__loadStockBasic()

    @staticmethod
    def getStockName(code):
        pd = StockBase().getStockBasic()
        row = pd.loc[pd['code'] == code, 'code_name']
        return "" if row.empty else row.values[0]

    @staticmethod
    def isTradingDate(date):
        return date in StockBase.getAllTradingDate()

    def __loadStockBasic(self):
        print("query_stock_basic ... ")
        bs = StockBase.getBaostock()
        rs = bs.query_stock_basic()
        data_list = []
        while (rs.error_code == '0') & rs.next():
            data_list.append(rs.get_row_data())
        result = pd.DataFrame(data_list, columns=rs.fields)
        print("size:" + str(result.shape))
        return result

    def __loadTradeDate(self):
        print("query_trade_dates ... ")
        bs = StockBase.getBaostock()
        rs = bs.query_trade_dates(start_date='2024-09-01', end_date=Objects.getDay(10))
        data_list = []
        while (rs.error_code == '0') & rs.next():
            tempDay = rs.get_row_data()
            if tempDay[1] == '1':
                data_list.append(tempDay[0])
        return data_list


if __name__ == '__main__':
    print(StockBase().getAllTradingDate())
    print(StockBase().getAllTradingDate())
    print(StockBase().isTradingDate('2020-01-02'))
    print(StockBase().isTradingDate('2020-01-02'))
    print(StockBase().getStockName('sh.000001'))
    print(StockBase().getStockName('sh.000006'))
    print(StockBase().getStockName('sh.000007'))
    # print(StockBase().getTradingDays('2024-01-01', '2025-02-01'))
    # print(StockBase().getTradingDays('2024-01-01', '2025-02-01'))
    # StockBase().exit()
